"""
VolatilityRegimeSwitch Strategy Implementation for Zipline
Adjusts portfolio exposure based on market volatility regimes
"""

from zipline.api import order_target_percent, symbol, get_datetime
from zipline.finance import commission, slippage
import pandas as pd

def initialize(context):
    # Strategy parameters
    context.vix_low = 20
    context.vix_high = 25
    context.volatility_cone_percentile = 0.9
    context.asset = symbol('SPY')  # Core asset
    
    # Leverage settings
    context.low_vol_leverage = 1.5  # 150% exposure
    context.high_vol_leverage = 0.5  # 50% exposure
    context.normal_leverage = 1.0   # 100% exposure
    
    # Risk management
    context.stop_loss_daily = 0.05  # 5% daily stop loss
    context.stop_loss_weekly = 0.10  # 10% weekly stop loss
    
    # Tracking variables
    context.last_vol_cone_update = None
    context.vol_cone = None
    context.current_leverage = context.normal_leverage
    context.last_week = None
    
    # Set commission and slippage
    context.set_commission(commission.PerShare(cost=0.005, min_trade_cost=1))
    context.set_slippage(slippage.FixedSlippage(spread=0.005))

def handle_data(context, data):
    current_date = get_datetime().normalize()
    
    # Initialize dates if not set
    if context.last_vol_cone_update is None:
        context.last_vol_cone_update = current_date
    if context.last_week is None:
        context.last_week = current_date
    
    # Monthly volatility cone update
    if (current_date - context.last_vol_cone_update).days >= 30:
        context.vol_cone = calculate_volatility_cone(context, data)
        context.last_vol_cone_update = current_date
    
    # Weekly return tracking
    if (current_date - context.last_week).days >= 7:
        context.portfolio.starting_weekly_value = context.portfolio.portfolio_value
        context.last_week = current_date
    
    # Daily volatility regime check
    check_vol_regime(context, data)
    
    # Daily stop loss check
    stop_loss_check(context, data)

def check_vol_regime(context, data):
    """Check volatility regime and adjust position accordingly"""
    # Get current VIX value
    current_vix = data.current(symbol('VIX'), 'close')
    
    # Get volatility cone value (if available)
    vol_cone = context.vol_cone if context.vol_cone else 0.5
    
    # Determine volatility regime
    if current_vix < context.vix_low and vol_cone < 0.25:
        # Low volatility regime
        set_leverage(context, data, context.low_vol_leverage)
    elif current_vix > context.vix_high and vol_cone > context.volatility_cone_percentile:
        # High volatility regime
        set_leverage(context, data, context.high_vol_leverage)
    else:
        # Normal volatility regime
        set_leverage(context, data, context.normal_leverage)

def set_leverage(context, data, leverage):
    """Adjust position to target leverage"""
    if leverage != context.current_leverage:
        order_target_percent(context.asset, leverage)
        context.current_leverage = leverage

def calculate_volatility_cone(context, data):
    """Calculate volatility cone (placeholder implementation)"""
    # In production, this would calculate historical volatility percentiles
    # This is a simplified placeholder
    return 0.6  # Example value

def stop_loss_check(context, data):
    """Implement daily and weekly stop loss logic"""
    portfolio = context.portfolio
    
    # Calculate daily returns
    daily_return = (portfolio.returns[-1] if len(portfolio.returns) > 0 else 0)
    
    # Calculate weekly returns
    weekly_return = (portfolio.portfolio_value / context.portfolio.starting_weekly_value) - 1
    
    # Implement stop loss rules
    if daily_return < -context.stop_loss_daily:
        reduce_risk_exposure(context, data)
    if weekly_return < -context.stop_loss_weekly:
        liquidate_positions(context, data)

def reduce_risk_exposure(context, data):
    """Reduce position size by 50%"""
    current_position = context.portfolio.positions[context.asset].amount
    if current_position > 0:
        target_position = current_position * 0.5
        order_target_percent(context.asset, target_position)

def liquidate_positions(context, data):
    """Liquidate all positions"""
    order_target_percent(context.asset, 0)